
#include <openrave-core.h>
#include <vector>
#include <iostream>

#include "orexample.h"
#include "pidcontroller.hpp"

using std::cout;
using std::endl;
using std::vector;
using std::string;

namespace cppexamples {

class RobotArm : public OpenRAVEExample {
 public:
  virtual void demothread(int argc, char ** argv) {
    penv->Load("models/world.xml");
    // attach a physics engine
    penv->SetPhysicsEngine(RaveCreatePhysicsEngine(penv,"ode"));
    penv->GetPhysicsEngine()->SetGravity(OpenRAVE::Vector(0,0,-9.8));

    // Find all the robots and get a pointer to them
    vector<OpenRAVE::RobotBasePtr> vrobots;
    penv->GetRobots(vrobots);
    OpenRAVE::RobotBasePtr probot = vrobots.at(0);

    vector<int> jointIndex;
    OpenRAVE::ControllerBasePtr jointController;

    { // Begin Mutually Exclusive Segment:
      OpenRAVE::EnvironmentMutex::scoped_lock lock(penv->GetMutex()); // lock environment
      jointController = RaveCreateController(penv,"odevelocity");
      vector<int> dofindices(probot->GetDOF());
      for(int i = 0; i < probot->GetDOF(); ++i) {
          dofindices[i] = i;
      }
      probot->SetController(jointController, dofindices,1) ;

      // set the velocity controller on all joints that have 'wheel' in their description
      for(vector<OpenRAVE::KinBody::JointPtr>::const_iterator itjoint = probot->GetJoints().begin(); itjoint != probot->GetJoints().end(); ++itjoint) {
        cout << "Got joint: " << (*itjoint)->GetName() << " - ";
        if( (*itjoint)->GetName().find("j") != string::npos ) {
          for(int i = 0; i < (*itjoint)->GetDOF(); ++i) {
            cout << "DOFIndex: " << (*itjoint)->GetDOFIndex()+i << endl;
            jointIndex.push_back((*itjoint)->GetDOFIndex()+i);
          }
        }
      }
    } // End Mutually Exclusive Segment

    // Attach controllers to the joints
    cout << "Attaching controllers.\n";
    cout << "jointIndex size: " << jointIndex.size() << endl;
    
    vector<PIDController*> controllers;
    for (unsigned int i=0; i<jointIndex.size(); i++) {
      PIDController* tmp = new PIDController();
      controllers.push_back(tmp);
    }

    vector<OpenRAVE::dReal> torques;          // joint torques
    torques.resize(jointIndex.size());
    
    vector<OpenRAVE::dReal> positions;        // angular position
    positions.resize(jointIndex.size());
    
    vector<OpenRAVE::dReal> setPoints;        // desired angle (set point)
    setPoints.resize(jointIndex.size());

    // Set points for joints 0, 1, 2:
    double j[3];
    j[0] = -30*M_PI/180;
    j[1] = 60*M_PI/180;
    j[2] = -90*M_PI/180;
    
    // PID Gains:
    double Kp = 200.0;
    double Ki = 0.7;
    double Kd = 1000;    // ~ 2*sqrt(Kp);
    
    // Max Joint Strength:
    double maxTorque = 1000; // Nm
    
    for(unsigned int i=0; i<jointIndex.size(); i++) {
      setPoints[i] = j[i%3];
      cout << "joint: " << i << " Kp: " << Kp << ", Ki: " << Ki << ", Kd: " << Kd << endl;
      controllers[i]->setKp(Kp);
      controllers[i]->setKi(Ki);
      controllers[i]->setKd(Kd);
      controllers[i]->setAntiWindup(-maxTorque, maxTorque);
    }

    cout << "Entering loop\n";
    cout << "Aiming for positions: " << setPoints[0] << ":" << setPoints[1] << ":" << setPoints[2] << "\n";

    uint64_t simtime = 0.0;
    //uint64_t old_time = 0.0;
    
    while(IsOk()) {
      {  // Begin Mutually Exclusive Segment:
        OpenRAVE::EnvironmentMutex::scoped_lock lock(penv->GetMutex()); // lock environment
        
        if( jointController ) {
          simtime = penv->GetSimulationTime()/1000;    // milliseconds
          cout << "time = " << simtime << "ms \n";
          // Get Joint positions:
          probot->GetDOFValues(positions);
          
          for(unsigned int i=0; i<jointIndex.size(); i++) {
            cout << "Joint:" << i << "  Pos:" << positions[i] << "  Set Point:" << setPoints[i] << "\n";
            
            // Get control signal from PID:
            double controlSignal = controllers[i]->compute(positions[i], setPoints[i]);
            
            // Max/Min to Actuator rails:
            double torque = controlSignal;
            if( controlSignal >  maxTorque ) torque =  maxTorque;
            if( controlSignal < -maxTorque ) torque = -maxTorque;
            
            torques[i] = torque;
            cout << "  Torque:" << torques[i] << "\n";
          }

          //old_time = simtime;

          // Set joint torques
          probot->SetDOFTorques(torques, false);
          //probot->SetDOFValues(setPoints, true);
        } 
        else {
          cout << "NO JOINT CONTROLLER OH NO\n";
        }
      }  // end Mutually Exclusive Segment
      
      // BUSY WAIT - TODO CHANGE THIS TO AN INTERRUPT:
      while( penv->GetSimulationTime()/1000 - simtime < 10 ) {   // wait 10 ms
        // wait...
      }
      //if( penv->GetSimulationTime()/1000 > 1000) exit(1); // stop sim after 1 second
    }
  }
};  // end class

} // end namespace cppexamples
